#!/usr/bin/env python3
"""
优化的FastAPI启动脚本
包含超时控制和资源管理
"""
import uvicorn
import signal
import sys
import os
import time
import torch
import gc

def cleanup_resources():
    """清理GPU资源"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()
        print("✅ GPU资源已清理")

def signal_handler(sig, frame):
    """信号处理器"""
    print("\n🛑 收到停止信号，正在清理资源...")
    cleanup_resources()
    sys.exit(0)

def main():
    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    print("🚀 启动优化版发音测评服务...")
    
    # 清理初始资源
    cleanup_resources()
    
    # 修复：设置正确的缓存路径
    os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 避免tokenizer警告
    os.environ["HF_HOME"] = "/root/.cache/huggingface"  # 修正：使用HF_HOME而不是过时的TRANSFORMERS_CACHE
    os.environ["TRANSFORMERS_OFFLINE"] = "0"  # 允许在线检查但优先使用缓存
    os.environ["HF_DATASETS_OFFLINE"] = "0"
    
    # 确保缓存目录存在
    cache_dir = "/root/.cache/huggingface"
    if not os.path.exists(cache_dir):
        os.makedirs(cache_dir, exist_ok=True)
        print(f"📁 创建缓存目录: {cache_dir}")
    else:
        print(f"📁 使用现有缓存目录: {cache_dir}")
    
    try:
        # 启动服务
        uvicorn.run(
            "fastapi_app:app",
            host="0.0.0.0",
            port=5001,
            workers=1,  # 单worker避免资源竞争
            timeout_keep_alive=30,
            log_level="info"
        )
    except Exception as e:
        print(f"❌ 服务启动失败: {e}")
        cleanup_resources()
        sys.exit(1)

if __name__ == "__main__":
    main() 